from dataset import train_loader,train_data
from model import Model
import torch


model =  Model(768)
model.change_model(train_data)
model =model.to("cuda:0")

#         optimizer.zero_grad()
#         batch = batch.to('cuda:0')
#         batch_size = batch['alarm'].batch_size
#         out = model(batch.x_dict, batch.edge_index_dict)
#         break
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
# batch = next(iter(train_loader))

for i in range(10):
    for batch in train_loader:
        optimizer.zero_grad()
        print("{}\n {}\n {}\n {}\n {}\n {}\n {}\n {}\n {}\n...................................................................".format(
              batch.x_dict["alarm"].shape,
              batch.x_dict["host"].shape,
              batch.x_dict["bussiness_tree"].shape,
              batch.edge_index_dict[('alarm', 'on', 'host')],
              batch.edge_index_dict[('host', 'rev_on', 'alarm')],
              batch.edge_index_dict[('alarm', 'to', 'bussiness_tree')],
              batch.edge_index_dict[('bussiness_tree', 'rev_to', 'alarm')],
              batch.edge_index_dict[('host', 'belongsto', 'bussiness_tree')],
              batch.edge_index_dict[('bussiness_tree', 'rev_belongsto', 'host')])
              )

        print("....................................................................")
        out1,out2,out3 = model(batch)
        aoh_edge_label = batch[("alarm", "on", "host")].edge_label
        atb_edge_label = batch[("alarm", "to", "bussiness_tree")].edge_label
        hbb_edge_label = batch[("host", "belongsto", "bussiness_tree")].edge_label
        print("out1!!!!!!!!!!!!!!!!!!!!!!!!!!!")
        print(out1.shape)
        print(out2.shape)
        print(out3.shape)
        print(aoh_edge_label.shape)
        print(atb_edge_label.shape)
        print(hbb_edge_label.shape)
        print("--------------------------------------------------------------------------------------------------")
        model(batch)



# data = data.to("cuda")